#!/usr/bin/env python3
"""
诊断CUDA问题的脚本
"""
import sys
import os

print("Step 1: Checking CUDA availability...")
print("=" * 80)

try:
    import torch
    print(f"✓ PyTorch imported successfully")
    print(f"  Version: {torch.__version__}")
    print(f"  CUDA available: {torch.cuda.is_available()}")

    if torch.cuda.is_available():
        print(f"  CUDA version: {torch.version.cuda}")
        print(f"  Device count: {torch.cuda.device_count()}")
        print(f"  Current device: {torch.cuda.current_device()}")
        print(f"  Device name: {torch.cuda.get_device_name(0)}")

        # Check for existing CUDA errors
        print("\nStep 2: Checking for existing CUDA errors...")
        try:
            torch.cuda.synchronize()
            print("✓ No existing CUDA errors")
        except Exception as e:
            print(f"✗ Existing CUDA error detected: {e}")
            print("  Recommendation: Restart Python kernel or reboot")
            sys.exit(1)

        # Test basic CUDA operation
        print("\nStep 3: Testing basic CUDA operations...")
        try:
            x = torch.randn(10, 10).cuda()
            y = x + x
            print("✓ Basic CUDA tensor operations work")
        except Exception as e:
            print(f"✗ Basic CUDA operations failed: {e}")
            sys.exit(1)

        # Test OverLock initialization WITHOUT CUDA
        print("\nStep 4: Testing OverLock on CPU...")
        try:
            from models_ext.overlock_local import overlock_t
            print("  Creating OverLock model on CPU...")
            model = overlock_t()
            print("✓ OverLock initialized on CPU successfully")

            # Check model structure
            total_params = sum(p.numel() for p in model.parameters())
            print(f"  Total parameters: {total_params:,}")

        except Exception as e:
            print(f"✗ OverLock initialization failed on CPU: {e}")
            print("  This suggests the problem is in OverLock code itself")
            import traceback
            traceback.print_exc()
            sys.exit(1)

        # Test moving to CUDA step by step
        print("\nStep 5: Testing OverLock move to CUDA...")
        print("  WARNING: This is where the error usually occurs")

        try:
            print("  Attempting to move model to CUDA...")
            model_cuda = model.to('cuda')
            torch.cuda.synchronize()  # Force synchronization
            print("✓ OverLock successfully moved to CUDA!")

            # Test forward pass with small input
            print("\nStep 6: Testing forward pass...")
            test_input = torch.randn(1, 3, 224, 224).cuda()
            with torch.no_grad():
                output = model_cuda.forward_features(test_input)
            print(f"✓ Forward pass successful!")
            print(f"  Output shape: {output[-1].shape}")

        except RuntimeError as e:
            print(f"✗ Moving to CUDA failed: {e}")
            print("\n" + "=" * 80)
            print("DIAGNOSIS:")
            print("=" * 80)

            if "assert" in str(e).lower():
                print("  Issue: CUDA device-side assertion")
                print("  Likely causes:")
                print("    1. OverLock has hardcoded assumptions about tensor sizes")
                print("    2. Pre-trained weights contain invalid values")
                print("    3. BatchNorm statistics are corrupted")
                print("\n  RECOMMENDED SOLUTIONS:")
                print("    A. Use a different backbone (ResNet, EfficientNet)")
                print("    B. Initialize OverLock without pretrained weights")
                print("    C. Check if OverLock version matches pretrained weights")

            import traceback
            traceback.print_exc()
            sys.exit(1)

    else:
        print("✗ CUDA is not available!")
        print("  Please check your CUDA installation")
        sys.exit(1)

except ImportError as e:
    print(f"✗ Import error: {e}")
    sys.exit(1)

print("\n" + "=" * 80)
print("✓ ALL TESTS PASSED!")
print("=" * 80)
print("\nConclusion: CUDA and OverLock are working correctly.")
print("The error must be occurring in a different context.")
